fix: linear_focused attention dimension mismatches in v_map_mlp and final_linear - #5
Merged
HirenMadhu merged 1 commit intoJul 24, 2026
Conversation
There was a problem hiding this comment.
Pull request overview
Fixes two shape/width mismatches in LorentzMultiheadAttention’s attention_type="linear_focused" code path so the value-side MLP and the head-concatenation projection (final_linear) align with the actual [..., 1:] (space-only) tensor widths used by linear-focused attention. This prevents runtime dimension errors while leaving the full attention path unchanged.
Changes:
- Adjust
v_map_mlpto map within the correct “v-space” width (out_channels - 2whenuse_weight=True, otherwiseout_channels - 1) to matchv = hyp_vs[..., 1:]andattn_output. - Make
final_linear’s input width conditional onlinear_focusedso it matches the concatenated linear-focused output width. - Fix the
reshapefeedingfinal_linearinlinear_focus_attentionto avoid hardcodingnum_heads * out_channelsand preserve the sequence length dimension.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Two dimension bugs in
LorentzMultiheadAttentionwhenattention_type="linear_focused":v_map_mlpmaps fromin_channels - 1toout_channels, but in the linear-focused path it receivesv = hyp_vs[..., 1:]whose last dimension isout_channels - 2(whenuse_weight=True). The result is added toattn_outputwhich also has dimensionout_channels - 2, so both the input and output ofv_map_mlpmust match that width.When
trans_heads_concat=True,final_linearinput width isnum_heads * out_channels, but the concatenated linear-focused output has widthnum_heads * (out_channels - 2). The reshape inlinear_focus_attentionalso hardcodes the wrong concat dimension.Both bugs cause a
RuntimeErroron any valid input to thelinear_focusedpath. Thefullattention path is unaffected.Scope
Code-only fix for the two dimension bugs I was running into. Does not change the linear-attention formulation itself.